# 前置参数：outputs，target_names
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
from matplotlib.colors import Normalize

# outputs = [[0, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 16],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 15, 6, 6, 16, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 6, 6, 6, 6],
#            [3, 3, 3, 3, 3, 1, 1, 2, 2, 5, 5, 6, 0, 0, 0, 0]]
# target_names = ['Alfalfa', 'Corn-notill', 'Corn-mintill', 'Corn', 'Grass-pasture', 'Grass-trees',
#                 'Grass-pasture-mowed', 'Hay-windrowed', 'Oats', 'Soybean-notill', 'Soybean-mintill',
#                 'Soybean-clean', 'Wheat', 'Woods', 'Buildings-Grass-Trees-Drives', 'Stone-Steel-Towers']
# 自定义颜色列表
category_colors = [
    '#89442c', '#0200fa', '#fc6500', '#02fd80', '#a24b9d', '#65b0fa',
    '#77fea7', '#395b6d', '#fffc05', '#feff7a', '#fc02fe', '#6401fb',
    '#00aff4', '#05fc05', '#aab153', '#60c03a'
]
# 创建一个颜色映射表，保留标签0为黑色
category_colors_with_black = ['#000000'] + category_colors  # 添加黑色作为标签0的颜色
# 创建一个颜色映射表
cmap = plt.cm.colors.ListedColormap(category_colors_with_black)
norm = Normalize(vmin=0, vmax=len(target_names))  # 归一化：确保vmin和vmax是基于类别数来设定

# 显示预测图像，使用 matplotlib 的 imshow 来绘制
plt.imshow(outputs, cmap=cmap, norm=norm)
plt.title('Predicted Image')

# 创建自定义图例
patches = [mpatches.Patch(color=category_colors_with_black[i], label=label)
           for i, label in enumerate(target_names)]

# 显示图例
plt.legend(handles=patches, bbox_to_anchor=(1.05, 0.5), loc='center left', borderaxespad=0.)

# 保存图像到本地文件
plt.savefig('predicted_image.png', bbox_inches='tight')  # 保存为PNG文件，'tight' 会自动调整边距

plt.show()